from tqdm import tqdm
import multiprocessing
import json
import openai
from argparse import ArgumentParser

# OpenAI全局配置
openai.base_url = "https://free.v36.cm/v1/"
openai.default_headers = {"x-foo": "true"}
# openai.api_key = "sk-ygs1G1n4rb74v8fJ51F6B4D3749a4d978aB99b4308D2A993"
openai.api_key = "sk-GbQ7V9AKYPUSLsu88e76A5C7B152449cAbFa0f5eCbD18aE3"

meta_prompt_qa = """
### Question
{text}

### Reference answer
{answer}

### Candidate answer
{predict}

Is the candidate answer correct based on the question and reference answer? 
Please only output a single "Yes" or "No". Do not output anything else.
""".strip()

def generate(item):
    try:
        # 构造prompt（保持原样）
        predict_value = item.get('predict', item.get('response', ''))
        prompt = meta_prompt_qa.replace('{text}', item['text'])\
                                .replace('{answer}', item['answer'])\
                                .replace('{predict}', predict_value)

        # 新版API调用方式
        completion = openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant who tries to help answer the user's question."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=0.5,
            top_p=0.95,
            n=3,
            max_tokens=1024,
            frequency_penalty=0,
            presence_penalty=0
        )
        
        # 提取结果
        rtn = [choice.message.content.strip() for choice in completion.choices]
        item['score'] = rtn
    except Exception as e:
        print(f"Error processing item {item.get('key', 'unknown')}: {str(e)}")
        item['score'] = ["Error"] * 3
    return item

def main():
    parser = ArgumentParser()
    parser.add_argument('--src_file', required=True)
    args = parser.parse_args()

    # 读取数据
    data = []
    with open(args.src_file, 'r') as f:
        for line in f:
            try:
                data.append(json.loads(line.strip()))
            except json.JSONDecodeError:
                print(f"Invalid JSON: {line}")

    # 多进程处理
    with multiprocessing.Pool(1) as pool:
        scores = list(tqdm(pool.imap(generate, data), total=len(data)))

    # 保存结果
    import os
    dir_path = os.path.dirname(args.src_file)
    base_name = os.path.basename(args.src_file)
    tgt_file = os.path.join(dir_path, f"result-{base_name}")
    # tgt_file = f"result-{args.src_file}"
    with open(tgt_file, "w", encoding='utf-8') as f:
        for d in scores:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")

if __name__ == '__main__':
    main()